/**
* file name  : it/ciano/jenoware/HWE.java
* authors    : Luciano Xumerle
* created    : mar 22 mar 2005 17:46:27 CET
*
* Copyright (c) 2004-2005 Luciano Xumerle. All rights reserved.
*
* This file is part of gephinfo.
*
* gephinfo is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program.  If not, see <http://www.gnu.org/licenses/>.
*
*/

package it.ciano.jenoware;

import java.util.*;
import it.ciano.util.*;

/**
 * Class HWE provide methods to compute chi-square test for deviations from Hardy-Weinberg
 * equilibrium in codominant systems.
 * The several constructors compares observed numbers of genotypes with the numbers of
 * genotypes expected under HWE, where the latter are computed on the basis
 * of allele frequencies estimated from the genotype frequencies (null hypothesis, H0).
 * For observed numbers, the relative cell frequencies are the estimates of the
 * genotype probabilities (alternative hypothesis, H1).
 * For the comparison between observed and expected numbers of genotypes,
 * chi-square is computed. If n is the number of alleles, the number of parameters
 * estimated under H0 and H1 is n-1 and n(n+1)/2-1, respectively.
 * Thus, the number of df for the chi-square test is n(n-1)/2.
 * The program also lists the log likelihoods (natural logarithms)
 * under H0 (no HWE) and H1  (HWE).
 * Results are unreliable with small numbers of observations.
 * The algorithm proposed is a port of the algorithm developed by J. Ott.
 * <br>
 * To improve the result of the test with small numbers of observations it's possible
 * to generate an exact test.
 * The p-value for the exact test (p-sim) is calculated by generating the possible samples
 * of genotype counts consistent with the observed allele counts, under a null
 * hypothesis of HW equilibrium, calculating the appropriate Chi-square test statistic
 * for each sample. The exact p-value is calculated by comparing the value of the
 * Chi-square test statistic of the observed data with the distribution of the test
 * statistic under the null hypothesis.
 * <br>References:<br>
 * http://www.meb.ki.se/genestat/tl/genass_ldmap/hwe_gterror/testing.htm
 *
 * @author Luciano Xumerle
 * @version 0.0.1
 */
public class HWE
{
    // total number of genotypes
    private int _totgeno;
    // store Chi-square value
    private double chisquare;
    // two tail test
    private double twotail;
    // store the degrees of freedom
    private double ndf;
    // store genotype Contributions to chi-square
    private double chi[];
    // observedAllele
    private int _observedAllele[];
    // observedGenotype
    private int _observedGenotype[];
    // expectedGenotypes
    private double _expectedGenotypes[];
    // number of simulations
    private int _stepNum;
    // array used to simulate the genotype variability
    private int _simulation[];
    // Hash to count the henotypes
    HashMap _counter;
    // p-sim
    private double psim;
    // random variable
    private Random _myRand;

    // likelihood Log likelihoods with(1)/without(2) HWE
    private double loglik1;
    private double loglik2;

    /**
     * Class constructor.
     * Genotypes must be furnished in the order,
     * 1/1, 1/2, 2/2, 1/3, 2/3, 3/3, 1/4, and so on.
     *
     * @param observedAllele array with the observed allele numbers.
     * @param observedGenotype array with the observed genotype numbers.
     */
    public HWE ( int[] observedAllele, int[] observedGenotype )
    {
        this ( observedAllele, observedGenotype, 0, false );
    }

    /**
     * Class constructor.
     * Genotypes must be furnished in the order,
     * 1/1, 1/2, 2/2, 1/3, 2/3, 3/3, 1/4, and so on.
     *
     * @param observedGenotype array with the observed genotype numbers.
     */
    public HWE ( int[] observedGenotype )
    {
        this ( null, observedGenotype, 0, false );
    }

    /**
     * Class constructor.
     * Genotypes must be furnished in the order,
     * 1/1, 1/2, 2/2, 1/3, 2/3, 3/3, 1/4, and so on.
     *
     * @param observedGenotype array with the observed genotype numbers.
     * @param stepNum number of simulations. If 0 then no simulations.
     * @param useNew if false then use an alternative method to randomize the genotypes.
     */
    public HWE ( int[] observedGenotype, int stepNum, boolean useNew )
    {
        this ( null, observedGenotype, stepNum, useNew );
    }

    /**
     * Class constructor.
     * Genotypes must be furnished in the order,
     * 1/1, 1/2, 2/2, 1/3, 2/3, 3/3, 1/4, and so on.
     *
     * @param observedAllele array with the observed allele numbers.
     * @param observedGenotype array with the observed genotype numbers.
     * @param stepNum number of simulations. If 0 then no simulations.
     * @param useNew if false then use an alternative method to randomize the genotypes.
     */
    public HWE ( int[] observedAllele, int[] observedGenotype, int stepNum, boolean useNew )
    {
        _stepNum = stepNum;
        _totgeno = 0;
        if ( observedAllele == null )
            _observedAllele = getAlleleFreqFromGenotypesFreq ( observedGenotype );
        else
            _observedAllele = observedAllele;
        _observedGenotype = observedGenotype;
        _expectedGenotypes = getExpectedGenotypes ( _observedAllele );
        chi = new double[ _expectedGenotypes.length ];
        int kk = 0;
        chisquare = 0;
        ndf = _observedAllele.length * ( _observedAllele.length - 1 ) / 2;
        for ( int i = 0; i < _observedAllele.length;i++ )
            for ( int j = 0; j <= i;j++ )
            {
                double tt = _observedGenotype[ kk ] - _expectedGenotypes[ kk ];
                tt = tt * tt;
                chi[ kk ] = tt / _expectedGenotypes[ kk ];
                chisquare += chi[ kk ];

                _totgeno += _observedGenotype[ kk++ ];
            }
        loglik1 = -_totgeno * Math.log ( _totgeno );
        loglik2 = loglik1;
        for ( int i = 0; i < _observedGenotype.length;i++ )
        {
            if ( _observedGenotype[ i ] > 0 )
                loglik1 += _observedGenotype[ i ] * Math.log ( _observedGenotype[ i ] );
            if ( _observedGenotype[ i ] > 0 )
                loglik2 += _observedGenotype[ i ] * Math.log ( _expectedGenotypes[ i ] );
        }
        twotail = SpecialFunction.chisqc ( ndf , chisquare );
        psim = 0;
        if ( _stepNum > 0 )
        {
            _myRand = new Random();
            psim = doSimulatedChiSquare ( stepNum, useNew );
        }
    }

    /**
     * Returns value of the calculated Chi-Square.
     *
     * @return the calculated Chi-Square value.
     */
    public double getChiSquare()
    {
        return chisquare;
    }

    /**
     * Returns an array with expected frequencies of the possible genotypes.
     * <br>Expected frequencies are calculated like <i>alleleFreq^2</i> for
     * homozigote genotype and <i>2 * alleleFreq1 * alleleFreq2</i> for
     * eterozigote genptype.
     *
     * @param observedAllele array with the number of counted alleles.
     * @return array with expected frequencies of Genotype.
     */
    public static double[] getExpectedGenotypes ( int[] observedAllele )
    {
        int alleleclasses = observedAllele.length;
        int totalAlleles = getTotalAlleles ( observedAllele );
        int totalGenotypes = totalAlleles / 2;
        int genotypeclasses = getGenotypeClassesFromAlleleClasses ( alleleclasses );
        double expectedGenotypes[] = new double[ genotypeclasses ];
        double alleleFreq[] = getAlleleFreq ( observedAllele );
        int kk = 0;
        for ( int i = 0; i < alleleclasses;i++ )
            for ( int j = 0; j <= i;j++ )
            {
                if ( i == j )
                    expectedGenotypes[ kk ] = alleleFreq[ i ] * alleleFreq[ i ];
                else
                    expectedGenotypes[ kk ] = 2 * alleleFreq[ i ] * alleleFreq[ j ];
                expectedGenotypes[ kk ] = totalGenotypes * expectedGenotypes[ kk ];
                kk++;
            }
        return expectedGenotypes;
    }

    /**
     * Returns the array with the frequencies for the observed alleles.
     * Array observedAllele is something like {50,30,60}.
     * Allele 1 has been observed 50 time, allele 2, 30 time and so on.
     *
     * @param observedAllele an integer array.
     * @return array with frequency.
     */
    public static double[] getAlleleFreq ( int[] observedAllele )
    {
        int totalAlleles = getTotalAlleles ( observedAllele );
        double alleleFreq[] = new double[ observedAllele.length ];
        for ( int i = 0; i < observedAllele.length; i++ )
            alleleFreq[ i ] = ( double ) observedAllele[ i ] / ( double ) totalAlleles;
        return alleleFreq;
    }

    /**
     * Returns the number of alleles observed.
     * Array observedAllele is something like {50,30,60}.
     * Allele 1 has been observed 50 time, allele 2, 30 time and so on.
     *
     * @param observedAllele an integer array.
     * @return the sum of all elements into observedAllele.
     */
    public static int getTotalAlleles ( int[] observedAllele )
    {
        int totalAlleles = 0;
        for ( int i = 0; i < observedAllele.length; i++ )
            totalAlleles += observedAllele[ i ];
        return totalAlleles;
    }

    /**
     * Overriding toString method.
     *
     * @return the String.
     */
    public String toString()
    {
        int ln = 6;
        String aa = "Allele number:\n";
        String nn = "Observed numbers of alleles:\n";
        String gg = "Genotypes:\n";
        String oo = "Observed numbers of genotypes:\n";
        String ee = "Expected numbers of genotypes:\n";
        String cc = "Contributions to chi-square:\n";
        String log = "\nLog likelihoods:\n"
                + CString.padLeft ( "No HWE:", 10, true )
                + CString.padLeft ( " " + loglik1, 15, true ) + "\n"
                + CString.padLeft ( "HWE:", 10, true )
                + CString.padLeft ( " " + loglik2, 15, true ) + "\n"
                + CString.padLeft ( "Diff:", 10, true )
                + CString.padLeft ( " " + ( loglik1 - loglik2 ), 15, true );

        int kk = 0;
        for ( int i = 0; i < _observedAllele.length; i++ )
        {
            aa += CString.padLeft ( " " + ( i + 1 ), ln, true );
            nn += CString.padLeft ( " " + _observedAllele[ i ] , ln, true );
            for ( int j = 0; j <= i;j++ )
            {
                gg += CString.padLeft ( " " + ( j + 1 ) + " " + ( i + 1 ), ln, true );
                oo += CString.padLeft ( " " + _observedGenotype[ kk ], ln, true );
                ee += CString.padLeft ( " " + _expectedGenotypes[ kk ], ln, true );
                cc += CString.padLeft ( " " + chi[ kk ], ln, true );
                kk++;
            }
        }

        String out = "Test for Hardy-Weinberg equilibrium\n"
                + "^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n"
                + aa + "\n"
                + nn + "\n\n"
                + gg + "\n"
                + oo + "\n"
                + ee + "\n"
                + cc + "\n"
                + "Chi-square = " + chisquare
                + "\tdf = " + ndf
                + "\t2p = " + twotail;
        if ( _stepNum > 0 )
            out += "\nNum of sim = " + _stepNum + "\tp-sim = " + psim;
        out += log;
        return out;
    }

    /**
     * Returns the p-sim.
     *
     * @param nstep number of simulation step.
     * @param useDefault true uses normal algorithm, false use Giovanni algorithm.
     * @return the p-sim.
     */
    private double doSimulatedChiSquare ( int nstep, boolean useDefault )
    {
        boolean dd = true;
        double pvalue = 0;
        for ( int i = 0; i < nstep; i++ )
        {
            initializeSimulation();
            double tt = 0;
            if ( useDefault )
            {
                tt = getSimulatedChiSquare ( dd );
                dd = ( dd && false );
            }
            else
                tt = getSimulatedChiSquareGio();
            if ( tt >= chisquare )
                pvalue++;
        }
        return pvalue / nstep;
    }

    /**
     * Returns the number of possible Genotype.
     *
     * @param alleleClasses number of observed Allele.
     * @return number of possible Genotype.
     */
    public static int getGenotypeClassesFromAlleleClasses ( int alleleClasses )
    {
        return ( alleleClasses * ( alleleClasses + 1 ) ) / 2;
    }

    /**
     * Returns number of observed alleles from number of observed genotypes.
     *
     * @param observedGenotype an array with number of observed genotypes.
     * @return an array with number of observed alleles.
     */
    public static int[] getAlleleFreqFromGenotypesFreq ( int[] observedGenotype )
    {
        int alleleclasses = getAlleleClassesFromGenotypeClasses ( observedGenotype.length );
        int res[] = new int[ alleleclasses ];
        for ( int i = 0; i < alleleclasses; i++ )
            res[ i ] = 0;
        int kk = 0;
        for ( int i = 0; i < alleleclasses;i++ )
            for ( int j = 0; j <= i;j++ )
            {
                res[ i ] += observedGenotype[ kk ];
                res[ j ] += observedGenotype[ kk ];
                kk++;
            }
        return res;
    }

    /**
     * Returns the number of allele classes from the number of genotype classes.
     *
     * @param genotypeClasses number of Genotype classes (i.e. 3, 6, 10, 15, ...).
     * @return number of allele classes (i.e. 2, 3, 4, 5, ...).
     */
    private static int getAlleleClassesFromGenotypeClasses ( int genotypeClasses )
    {
        return ( int ) ( Math.sqrt ( 1 + 8 * genotypeClasses ) - 1 ) / 2;
    }

    /**
     * Returns the value of a simulated Chi-Square Test.
     * This method uses a randomize method to generate a group of Genotype with
     * the same frequencies allele observed into the original population.
     *
     * @param fromZero true start from the beginning of array, false start from tail.
     * @return the simulated value.
     */
    private double getSimulatedChiSquare ( boolean fromZero )
    {
        _counter = new HashMap();
        int lastElement = _simulation.length - 1;
        // swap array elements and count resulting genotypes
        if ( fromZero )
        {
            for ( int i = 0; i < _simulation.length; i += 2 )
            {
                swapElements ( i, i + _myRand.nextInt ( lastElement - i ) );
                updateCounterHashMap ( _simulation[ i ], _simulation[ i + 1 ] );
            }
        }
        else
        {
            for ( int i = lastElement; i > 0; i -= 2 )
            {
                swapElements ( i, _myRand.nextInt ( i ) );
                updateCounterHashMap ( _simulation[ i ], _simulation[ i - 1 ] );
            }
        }
        return new HWE ( _observedAllele,
                getObservedGenotypeFromHashMap() ).getChiSquare();
    }

    /**
     * Returns the value of a simulated Chi-Square Test.
     * This method uses a randomize method to generate a group of Genotype with
     * the same frequencies allele observed into the original population.
     * Randomize method by Giovanni Malerba (more CPU intesitive).
     *
     * @return the simulated value.
     */
    private double getSimulatedChiSquareGio()
    {
        _counter = new HashMap();
        int lastElement = _simulation.length - 1;
        // swap array elements and count resulting genotypes
        for ( int i = 0; i < _simulation.length; i++ )
        {
            swapElements ( i, _myRand.nextInt ( lastElement ) );
        }
        // count genotypes
        for ( int i = 0; i < _simulation.length; i += 2 )
            updateCounterHashMap ( _simulation[ i ], _simulation[ i + 1 ] );
        return new HWE ( _observedAllele,
                getObservedGenotypeFromHashMap() ).getChiSquare();
    }

    /**
     * Parses the HashMap _counter and returns the observed Genotype array.
     *
     * @return the observed Genotype array.
     */
    private int[] getObservedGenotypeFromHashMap()
    {
        if ( _counter == null )
            return null;
        int obsGeno[] = new int[ _observedGenotype.length ];
        int kk = 0;
        for ( int i = 1; i <= _observedAllele.length;i++ )
            for ( int j = 1; j <= i;j++ )
            {
                int p = 0;
                String temp = j + ";" + i;
                if ( _counter.containsKey ( temp ) )
                    p = ( ( Integer ) _counter.get ( temp ) ).intValue();
                obsGeno[ kk++ ] = p;
            }
        return obsGeno;
    }

    /**
     * Updates the HashMap _counter with the genotype given by the two integer.
     *
     * @param a first element.
     * @param b second element.
     */
    private void updateCounterHashMap ( int a, int b )
    {
        StringBuffer strBuf = new StringBuffer();
        if ( a > b )
        {
            strBuf.append ( b );
            strBuf.append ( ";" );
            strBuf.append ( a );
        }
        else
        {
            strBuf.append ( a );
            strBuf.append ( ";" );
            strBuf.append ( b );
        }
        int p;
        String temp = strBuf.toString();
        if ( !_counter.containsKey ( temp ) )
            p = 1;
        else
            p = ( ( Integer ) _counter.get ( temp ) ).intValue() + 1;
        _counter.put ( temp, new Integer ( p ) );
    }

    /**
     * Initialize simulation array.
     */
    private void initializeSimulation()
    {
        int total = getTotalAlleles ( _observedAllele );
        _simulation = new int[ total ];
        int kk = 0;
        int allele = 0;
        for ( int i = 0; i < _observedAllele.length; i++ )
        {
            allele++;
            for ( int j = 0; j < _observedAllele[ i ]; j++ )
                _simulation[ kk++ ] = allele;
        }
    }

    /**
     * Swaps two elements into the simulation array.
     *
     * @param a first element index.
     * @param b second element index.
     */
    private void swapElements ( int a, int b )
    {
        int temp = _simulation[ a ];
        _simulation[ a ] = _simulation[ b ];
        _simulation[ b ] = temp;
    }


} // end class

